import lightgbm as lgb
from sklearn.metrics import roc_auc_score
import numpy as np


def sigmoid(x):
    return 1. / (1 + np.exp(-x))


class LgbModel:
    def __init__(self, feature_names=None, category_features=None):
        self.model_name = "lgb"
        self.feature_names = feature_names
        if category_features:
            self.category_features = category_features
        else:
            self.category_features = []
        from multiprocessing import cpu_count

        self.model = lgb.LGBMClassifier(
            boosting_type='gbdt',
            # boosting_type='dart',
            num_leaves=10, reg_alpha=1, reg_lambda=3, use_missing=True,
            max_depth=8, n_estimators=100, objective='binary',
            subsample=0.8, colsample_bytree=0.6, subsample_freq=1,
            # is_unbalance=False,
            scale_pos_weight=1,
            learning_rate=0.5, min_child_weight=60, n_jobs=cpu_count() - 1,
            random_state=2022, importance_type='gain'  # split(use num)/gain
        )

    def fit_(self, train_x, train_y, test_x, test_y, verbose=True):
        self.model.fit(train_x, train_y, eval_set=[(train_x, train_y), (test_x, test_y)],
                       feature_name=self.feature_names, categorical_feature=self.category_features,
                       early_stopping_rounds=100,
                       eval_metric='auc', verbose=verbose)
        # clf.evals_result_
        return self.model, self.model.best_score_['training'], self.model.best_score_['valid_1']

    def model_predict(self, eval_data_x, raw_score=True, pred_leaf=False, pred_contrib=False):
        return self.model.predict(eval_data_x, raw_score=raw_score, pred_leaf=pred_leaf, pred_contrib=pred_contrib)

    def score(self, eval_x, eval_y):
        predict_res = self.model_predict(eval_x)
        predict_sigmoid = sigmoid(predict_res)
        eval_auc = roc_auc_score(eval_y, predict_sigmoid)
        return eval_auc

    def get_model_params(self, deep=False):
        return self.model.get_params(deep=deep)

    def save_model(self, save_path, model_prefix=""):
        return self.model.booster_.save_model(save_path + f"{model_prefix}_{self.model_name}_model.txt")

    @property
    def best_iter_num(self):
        return self.model.best_iteration_

    @property
    def get_model_name(self):
        return self.model_name

    @property
    def get_feat_importance(self):
        return self.model.feature_importances_

    @property
    def get_model_info(self):
        return self.model